#!/usr/bin/env python3
import os
import sys
import argparse
from pathlib import Path
from mistralai import Mistral
import json
import base64
import re
import shutil

def save_image_from_base64(base64_str, output_dir, image_name):
    """Save a base64 encoded image to a file."""
    try:
        # Extract the actual base64 content if it includes the data URI scheme
        if "," in base64_str:
            base64_content = base64_str.split(",", 1)[1]
        else:
            base64_content = base64_str
            
        # Decode the base64 string to binary data
        image_data = base64.b64decode(base64_content)
        
        # Create the output directory if it doesn't exist
        output_dir.mkdir(exist_ok=True, parents=True)
        
        # Save the image to a file
        output_path = output_dir / image_name
        with open(output_path, "wb") as image_file:
            image_file.write(image_data)
            
        return output_path
    except Exception as e:
        print(f"Error saving image {image_name}: {e}")
        return None

def extract_markdown_from_pdf(pdf_path, client, image_dir):
    """Extract markdown from a PDF file using Mistral OCR API."""
    try:
        print(f"Uploading {pdf_path} to Mistral...")
        
        # Upload file directly to Mistral with correct format
        uploaded_file = client.files.upload(
            file={
                "file_name": pdf_path.name,
                "content": pdf_path.read_bytes(),
            },
            purpose="ocr"
        )
        
        # Get a signed URL for the uploaded file
        signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1)
        
        print(f"Processing the PDF with OCR...")
        ocr_response = client.ocr.process(
            model="mistral-ocr-latest",
            document={
                "type": "document_url",
                "document_url": signed_url.url
            },
            include_image_base64=True
        )
        
        # Create a folder for this specific PDF's images
        pdf_image_dir = image_dir / pdf_path.stem
        pdf_image_dir.mkdir(exist_ok=True, parents=True)
        
        # Process and combine markdown from all pages
        all_markdown = ""
        for page_idx, page in enumerate(ocr_response.pages):
            # Handle images in the page
            page_markdown = page.markdown
            for img_idx, img in enumerate(page.images):
                # Create a unique filename for the image
                image_filename = f"page{page_idx+1}_img{img_idx+1}{get_extension_from_base64(img.image_base64)}"
                
                # Save the image to a file
                save_image_from_base64(
                    img.image_base64, 
                    pdf_image_dir, 
                    image_filename
                )
                
                # Create a relative path for markdown
                relative_image_path = f"pic/{pdf_path.stem}/{image_filename}"
                
                # Replace the base64 image reference with a link to the local file
                page_markdown = page_markdown.replace(
                    f"![{img.id}]({img.id})",
                    f"![{img.id}]({relative_image_path})"
                )
                
                # Also replace any direct references to the base64 data
                if img.image_base64 in page_markdown:
                    page_markdown = page_markdown.replace(
                        f"![{img.id}]({img.image_base64})",
                        f"![{img.id}]({relative_image_path})"
                    )
                
            all_markdown += page_markdown + "\n\n"
            
        return all_markdown
    except Exception as e:
        print(f"Error extracting markdown from {pdf_path}: {e}")
        return None

def get_extension_from_base64(base64_str):
    """Determine the file extension based on the base64 data."""
    # Check if it's a data URI with MIME type
    if base64_str.startswith('data:'):
        mime_match = re.match(r'data:(image/\w+);base64,', base64_str)
        if mime_match:
            mime_type = mime_match.group(1)
            if mime_type == 'image/jpeg' or mime_type == 'image/jpg':
                return '.jpg'
            elif mime_type == 'image/png':
                return '.png'
            elif mime_type == 'image/gif':
                return '.gif'
            elif mime_type == 'image/tiff':
                return '.tiff'
            elif mime_type == 'image/bmp':
                return '.bmp'
            elif mime_type == 'image/webp':
                return '.webp'
            elif mime_type == 'image/svg+xml':
                return '.svg'
    
    # Default to png if we can't determine the type
    return '.png'

def process_markdown_file(md_path, image_dir):
    """Process a markdown file to extract base64 images and save them locally."""
    print(f"Processing markdown file: {md_path}")
    
    # Read the markdown file
    with open(md_path, "r", encoding="utf-8") as file:
        content = file.read()
    
    # Define the folder name for extracted images
    folder_name = md_path.stem
    image_folder = image_dir / folder_name
    image_folder.mkdir(exist_ok=True, parents=True)
    
    # Find base64 image patterns in markdown
    # Look for both ![...](data:image/...) and regular data:image URLs
    image_pattern = r'!\[(.*?)\]\((data:image/[^;]+;base64,[^)]+)\)'
    url_pattern = r'(data:image/[^;]+;base64,[^"\'\s]+)'
    
    # Replace inline base64 images with local file references
    img_count = 0
    
    # Process images in markdown image tags
    for match in re.finditer(image_pattern, content):
        img_count += 1
        alt_text = match.group(1)
        base64_data = match.group(2)
        
        # Get file extension from base64 data
        extension = get_extension_from_base64(base64_data)
        image_filename = f"img{img_count}{extension}"
        
        # Save the image
        save_image_from_base64(base64_data, image_folder, image_filename)
        
        # Create relative path for the image
        relative_path = f"pic/{folder_name}/{image_filename}"
        
        # Replace the base64 data with the local file reference
        content = content.replace(match.group(0), f"![{alt_text}]({relative_path})")
    
    # Process any remaining standalone base64 image URLs
    # This handles cases where the base64 data might be in other contexts
    for match in re.finditer(url_pattern, content):
        if match.group(0) not in content:  # Skip if already replaced
            continue
            
        img_count += 1
        base64_data = match.group(1)
        
        # Get file extension from base64 data
        extension = get_extension_from_base64(base64_data)
        image_filename = f"img{img_count}{extension}"
        
        # Save the image
        save_image_from_base64(base64_data, image_folder, image_filename)
        
        # Create relative path for the image
        relative_path = f"pic/{folder_name}/{image_filename}"
        
        # Replace the base64 data with the local file reference
        content = content.replace(match.group(0), relative_path)
    
    # Write the updated markdown back to the file
    with open(md_path, "w", encoding="utf-8") as file:
        file.write(content)
    
    print(f"Processed {img_count} images from {md_path}")
    return img_count > 0

def process_pdf_file(pdf_path, text_dir, image_dir, client):
    """Process a single PDF file and convert it to markdown."""
    print(f"Processing {pdf_path}")
    
    # Get base filename without extension
    base_name = pdf_path.stem
    
    # Define output markdown file path
    md_path = text_dir / f"{base_name}.md"
    
    # Check if markdown file already exists
    if md_path.exists():
        print(f"Skipping {pdf_path} - {md_path} already exists")
        return True
    
    # Extract markdown from PDF using Mistral OCR
    extracted_markdown = extract_markdown_from_pdf(pdf_path, client, image_dir)
    
    if extracted_markdown:
        # Write markdown to output file
        try:
            with open(md_path, "w", encoding="utf-8") as md_file:
                md_file.write(extracted_markdown)
            print(f"Successfully converted {pdf_path} to {md_path}")
            return True
        except Exception as e:
            print(f"Error writing to {md_path}: {e}")
            return False
    else:
        print(f"Failed to extract markdown from {pdf_path}")
        return False

def main():
    # Set up argument parser
    parser = argparse.ArgumentParser(description='Convert PDF to markdown using Mistral OCR API')
    
    # Create a mutually exclusive group for input options
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument('input_path', type=str, nargs='?', help='PDF file or directory containing PDF files to process')
    input_group.add_argument('--file', type=str, help='Process a specific PDF file')
    input_group.add_argument('--file-md', type=str, help='Process a markdown file with embedded base64 images')
    
    parser.add_argument('--force', action='store_true', help='Process files even if markdown already exists')
    parser.add_argument('--image-dir', type=str, default='pic', help='Directory to save images (default: pic)')
    args = parser.parse_args()
    
    # Define directories
    image_dir = Path(args.image_dir)
    
    # Create image directory if it doesn't exist
    image_dir.mkdir(exist_ok=True)
    
    # Option to process an existing markdown file with embedded base64 images
    if args.file_md:
        md_file = Path(args.file_md)
        if not md_file.exists():
            print(f"Error: Markdown file {md_file} not found")
            sys.exit(1)
            
        if md_file.suffix.lower() != '.md':
            print(f"Error: File {md_file} is not a markdown file")
            sys.exit(1)
            
        process_markdown_file(md_file, image_dir)
        print(f"Completed processing markdown file: {md_file}")
        sys.exit(0)
        
    # For PDF processing, need Mistral API key
    api_key = os.environ.get("MISTRAL_API_KEY")
    if not api_key:
        print("Error: MISTRAL_API_KEY environment variable not set")
        print("Set it with: export MISTRAL_API_KEY=your_api_key")
        sys.exit(1)
        
    # Initialize Mistral client
    client = Mistral(api_key=api_key)
    
    # Define directories for PDF processing
    papers_dir = Path("papers")
    text_dir = Path("text")
    
    # Create text directory if it doesn't exist
    text_dir.mkdir(exist_ok=True)
    
    # Process specific file provided via --file option
    if args.file:
        specific_file = Path(args.file)
        if not specific_file.exists():
            print(f"Error: File {specific_file} not found")
            sys.exit(1)
            
        if not specific_file.suffix.lower() == '.pdf':
            print(f"Error: File {specific_file} is not a PDF")
            sys.exit(1)
            
        # Check if output file exists and --force is not set
        output_file = text_dir / f"{specific_file.stem}.md"
        if output_file.exists() and not args.force:
            print(f"Skipping {specific_file} - {output_file} already exists")
            print("Use --force to process anyway")
            sys.exit(0)
            
        # If force is set and output exists, remove the existing image directory
        if args.force and output_file.exists():
            pdf_image_dir = image_dir / specific_file.stem
            if pdf_image_dir.exists():
                print(f"Removing existing images in {pdf_image_dir}")
                shutil.rmtree(pdf_image_dir, ignore_errors=True)
            
        process_pdf_file(specific_file, text_dir, image_dir, client)
    # Process file or directory provided as positional argument
    elif args.input_path:
        input_path = Path(args.input_path)
        
        # Check if the input path exists
        if not input_path.exists():
            print(f"Error: {input_path} not found")
            sys.exit(1)
        
        # If it's a file, process it
        if input_path.is_file():
            if not input_path.suffix.lower() == '.pdf':
                print(f"Error: File {input_path} is not a PDF")
                sys.exit(1)
                
            # Check if output file exists and --force is not set
            output_file = text_dir / f"{input_path.stem}.md"
            if output_file.exists() and not args.force:
                print(f"Skipping {input_path} - {output_file} already exists")
                print("Use --force to process anyway")
                sys.exit(0)
                
            # If force is set and output exists, remove the existing image directory
            if args.force and output_file.exists():
                pdf_image_dir = image_dir / input_path.stem
                if pdf_image_dir.exists():
                    print(f"Removing existing images in {pdf_image_dir}")
                    shutil.rmtree(pdf_image_dir, ignore_errors=True)
                
            process_pdf_file(input_path, text_dir, image_dir, client)
        # If it's a directory, process all PDFs in it
        elif input_path.is_dir():
            pdf_files = list(input_path.glob("*.pdf"))
            
            if not pdf_files:
                print(f"No PDF files found in {input_path}")
                sys.exit(0)
            
            print(f"Found {len(pdf_files)} PDF files to process in {input_path}")
            
            # Count skipped files
            processed = 0
            skipped = 0
            failed = 0
            
            for pdf_path in pdf_files:
                # Check if output file exists and --force is not set
                output_file = text_dir / f"{pdf_path.stem}.md"
                if output_file.exists() and not args.force:
                    print(f"Skipping {pdf_path} - {output_file} already exists")
                    skipped += 1
                    continue
                    
                # If force is set and output exists, remove the existing image directory
                if args.force and output_file.exists():
                    pdf_image_dir = image_dir / pdf_path.stem
                    if pdf_image_dir.exists():
                        print(f"Removing existing images in {pdf_image_dir}")
                        shutil.rmtree(pdf_image_dir, ignore_errors=True)
                    
                # Process the file
                if process_pdf_file(pdf_path, text_dir, image_dir, client):
                    processed += 1
                else:
                    failed += 1
            
            # Print summary
            print(f"\nSummary:")
            print(f"  Total files: {len(pdf_files)}")
            print(f"  Successfully processed: {processed}")
            print(f"  Skipped (already exist): {skipped}")
            print(f"  Failed: {failed}")
    else:
        # Default behavior: process all PDF files in papers directory
        # Check if papers directory exists
        if not papers_dir.exists():
            print(f"Error: {papers_dir} directory not found")
            sys.exit(1)
        
        # Process all PDF files in papers directory
        pdf_files = list(papers_dir.glob("*.pdf"))
        
        if not pdf_files:
            print(f"No PDF files found in {papers_dir}")
            sys.exit(0)
        
        print(f"Found {len(pdf_files)} PDF files to process in {papers_dir}")
        
        # Count skipped files
        processed = 0
        skipped = 0
        failed = 0
        
        for pdf_path in pdf_files:
            # Check if output file exists and --force is not set
            output_file = text_dir / f"{pdf_path.stem}.md"
            if output_file.exists() and not args.force:
                print(f"Skipping {pdf_path} - {output_file} already exists")
                skipped += 1
                continue
                
            # If force is set and output exists, remove the existing image directory
            if args.force and output_file.exists():
                pdf_image_dir = image_dir / pdf_path.stem
                if pdf_image_dir.exists():
                    print(f"Removing existing images in {pdf_image_dir}")
                    shutil.rmtree(pdf_image_dir, ignore_errors=True)
                
            # Process the file
            if process_pdf_file(pdf_path, text_dir, image_dir, client):
                processed += 1
            else:
                failed += 1
        
        # Print summary
        print(f"\nSummary:")
        print(f"  Total files: {len(pdf_files)}")
        print(f"  Successfully processed: {processed}")
        print(f"  Skipped (already exist): {skipped}")
        print(f"  Failed: {failed}")

if __name__ == "__main__":
    main() 